"""

    File: decoder_py.py
    Description: 
    
        This module decomposes a raw RTMP stream into AMF packets.  It keeps a buffer of incoming data for each 
        stream and drops complete packets into a buffer.

    Author: Kyle Vogt
    Date  : June 5th, 2007
    Copyright (c) 2007, Justin.tv, Inc.
    
"""

import struct, traceback
from twisted.python import log
from pms.rtmp import amf_py as amf
import os, traceback

class Decoder():
    
    READ_HEADER = 0x01
    READ_BODY = 0x02
    
    def __init__(self, max_buffer = 5000000):
        self.amf = amf.Amf()
        self.data = ''
        self.pointer = 0
        self.max_buffer = max_buffer
        self.stream = None
        self.streams = {}
        self.packets = []
        self.events = []
        self.bytesRead = 3073
        self.chunkSize = 128
        self.count = 0
        self.validPackets = [0x08, 0x09, 0x12, 0x14, 0x16]
        self.statusPackets = [0x04]
        self.eventPackets = [0x12, 0x14]
        self.mode = self.READ_HEADER
        self.left_in_chunk = 0
                
    def process(self, data):
        "Process raw RTMP data and parse it into AMF streams"
        
        # Keep track of total bytes read
        self.data += data
        self.bytesRead += len(data)
        
        # Check to see if data is corrupt and filling up our buffer
        if len(self.data) > self.max_buffer:
            log.msg("Error: decoder scanned %i bytes without finding a packet" % self.max_buffer)
            self.data = ''
            return False
        
        while True:
            self.available_bytes = len(self.data) - self.pointer
            # Do we need to read the header?
            if self.mode == self.READ_HEADER:
                # Can we read header length?
                if self.available_bytes >= 1:
                    self.headerLength = self.readHeaderLength()
                else:
                    return True
                # Can we read the full header?
                if self.available_bytes >= self.headerLength:
                    self.stream = self.readHeader()
                    # Did we actually not have enough bytes? (special timecode)
                    if not self.stream: 
                        return True
                    # All good, read the body now
                    else:
                        self.pointer += self.headerLength
                        self.available_bytes -= self.headerLength
                        self.mode = self.READ_BODY
                # Can't read whole header, we're done for now
                else:
                    return True
            # Do we need to read the body?
            if self.mode == self.READ_BODY:
                # How much is left to go?
                remaining = self.stream['length'] - len(self.stream['data'])
                bytes_to_read = min(self.chunkSize, remaining)
                # Do we have enough bytes right now?
                if self.available_bytes < bytes_to_read:
                    return True
                # Read up to an entire chunk
                self.stream['data'] += self.data[self.pointer : self.pointer + bytes_to_read]
                self.pointer += bytes_to_read
                self.available_bytes -= bytes_to_read
                remaining -= bytes_to_read
                # Is the packet complete?
                if not remaining:
                    self.processPacket()
                    self.mode = self.READ_HEADER
                # If we read a full chunk, it's time to read a 1-byte header
                if bytes_to_read == self.chunkSize:
                    self.mode = self.READ_HEADER
                # Are we out of data to work with?
                if not self.available_bytes:
                    self.data = ''
                    self.pointer = 0
                    return
                                
    def readHeader(self):
        # Read the stream id
        objectId = str(self.readObjectId())
        if objectId in self.streams:
            stream = self.streams[objectId]
        else:
            stream = self.createStream(objectId)
        # Analyze header
        if self.headerLength > 1: 
            stream['timecode'] = self.readTimecode()    
            stream['absflag'] = False
        if self.headerLength > 4: 
            stream['type'] = self.readType()
            stream['length'] = self.readLength()
        if self.headerLength > 8: 
            stream['id'] = self.readStreamId()
            stream['absflag'] = True
        # Handle special timecode case
        if stream['timecode'] == 16777215:
            log.msg('Long timecode detected!')
            self.headerLength += 4
            if self.available_bytes < self.headerLength: 
                return False
            stream['timecode'] == self.readSpecialTimecode()
        return stream
        
    def processPacket(self):
        "Process a complete packet and drop it into the packet buffer"
        result = True
        #print "got a packet of type %s" % self.stream['type']
        if self.stream['type'] == 0x01: 
            if len(self.stream['data']) >= 4:
                newSize = self.readChunkSize(self.stream['data'])
                if newSize >= 64 and newSize <= 65536: 
                    self.chunkSize = newSize
                    log.msg("Chunk size set to %i" % newSize)
                else: 
                    log.msg("Error: decoder found corrupt data (tried to set chunk size to %s)" % newSize)
        # Useful packes
        elif self.stream['type'] in self.validPackets:
            packet = {'count' : self.count}
            for key in self.stream.keys(): packet[key] = self.stream[key]
            packet['key'] = (packet['type'] == 0x09 and self.isKeyframe(packet))
            self.count += 1
            # Event packet
            if self.stream['type'] in self.eventPackets:
                event = {'id' : self.stream['id'], 'data' : self.amf.decode(self.stream['data'])}
                #log.msg('event: %s' % event)
                self.events.append(event)
            # Regular A/V packet
            else:
                self.packets.append(packet)
        elif self.stream['type'] in self.statusPackets:
            pass
        # Random packet we don't care about
        else:
            pass
        # Clean up for next packet
        self.stream['data'] = ''
            
    def createStream(self, objectId):
        "Create and return new stream buffer"
        #log.msg('Created stream %s' % objectId)
        if int(objectId) > 16:
            log.msg('May have corrupt data - netstream id %s was just created' % objectId)
        stream = {'data' : '', 'type' : 0, 'length' : 0, 'absflag' : False, 'timecode' : 0, 'id' : 0, 'object' : int(objectId)}
        self.streams[objectId] = stream
        return stream
                        
    def readChunkSize(self, data):
        "Return the new chunk size for data blocks"
        size = struct.unpack("!i", data[:4])[0]    
        return size
                        
    def readHeaderLength(self):
        "Return header length in bytes (2 LSB's)"
        try:
            bits = ord(self.data[self.pointer]) >> 6
        except:
            log.msg('Error reading header length!')
            traceback.print_exc()
            log.msg('Pointer: %s len(data): %s' % (self.pointer, len(self.data)))
            return 0
        if bits == 0: return 12
        elif bits == 1: return 8
        elif bits == 2: return 4
        elif bits == 3: return 1
        else: return 0
        
    def readObjectId(self):
        "Return rtmp object id (mask 0x3F)"
        return (ord(self.data[self.pointer]) & 0x3F)        
        
    def readStreamId(self):
        "Read rtmp stream id"
        return struct.unpack("<I", self.data[self.pointer + 8 : self.pointer + 12])[0]

    def readLength(self):
        "Read rtmp packet length"
        return struct.unpack("!I", chr(0) + self.data[self.pointer + 4: self.pointer + 7])[0]

    def readTimecode(self):
        "Read rtmp packet timecode"
        return struct.unpack("!I", chr(0) + self.data[self.pointer + 1 : self.pointer + 4])[0]

    def readSpecialTimecode(self):
        "Read rtmp packet timecode"
        return struct.unpack("!I", self.data[self.pointer + self.headerLength - 4 : self.pointer + self.headerLength])[0]

    def readType(self):
        "Read rtmp packet type"
        return ord(self.data[self.pointer + 7])
        
    def isKeyframe(self, tag):
        "Boolean fucntion to determine whether the FLV tag in this tag is a keyframe"
        return ((len(tag['data']) > 0) and (tag['type'] == 0x09) and self.getVideoFrame(tag) == 'keyframe')

    def getVideoFrame(self, tag):
        "Determines whether the FLV tag is a keyframe or not"
        frame = (ord(tag['data'][0]) & 0xF0) >> 4
        if frame == 1: return 'keyframe'
        elif frame == 2: return 'iframe'
        elif frame == 3: return 'unsupported'
        else: return 'err'    
